{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Estimating the number of classes in a population"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A notebook to evalute various estimators of the number of distinct classes in a population. \n",
    "\n",
    "Intended to supplement the readings *On the Estimation of the Number of Classes in a Population* by Leo A Goodman and *The Population Frequencies of Species and the Estimation of Population Parameters* by I. J. Good"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Notebook structure\n",
    "\n",
    "- #### Section 1\n",
    " - 1.1 Overview Goodman's notation & explain the relationship between Goodman's problem statement in 1949 and a Twitter data problem\n",
    "\n",
    "- #### Section 2\n",
    " - 2.1 Make a test population with class frequencies that follow a) Poisson distribution or b) uniform distribution\n",
    " - 2.2 Calculate the unbiased estimator suggested in Goodman's paper\n",
    " - 2.3 Calculate and discuss other (biased) estimators suggested by Goodman\n",
    "\n",
    "- #### Section 3\n",
    " - 3.1 Load real data from a Twitter dataset (all Tweets corresponding to some PowerTrack query)\n",
    " - 3.2 Pull a random sample of size n of this data\n",
    " - 3.3 Calculate the estimators again for Twitter data\n",
    " - 3.4 Exploration and suggestions for future work"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Section 1: Notation & problem statement\n",
    "\n",
    "#####Goodman's problem: \n",
    "The 1949 paper refers to \"estimating the number of classes in a population.\" Imagine a population of a large number (*N*) of differently colored balls in an urn. Given some random draw of *n* balls from the urn, from the distribution of those colors estimate the total number of different colors (classes, *K*) in the urn. \n",
    "\n",
    "##### Social data:\n",
    "The problem can be extended to social data. Given that there are *N* total Tweets about cats, select a much smaller (and more manageable) sample of *n* cat-related Tweets. Call the user who created a Tweet the \"class\" of that Tweet. Now use the distribution of Tweets per user in your *n*-sized sample to estimate the total number of unique users Tweeting about cats. I'll continue to refer to the problem this way.\n",
    "\n",
    "#####Notation:\n",
    "* *__N__*: The total size of the population, or the total number of cat Tweets on Twitter\n",
    "* *__n__*: The size of the sample, the number of cat Tweets that we are actually going to inspect\n",
    " * We can think about n as a vector with *K* elements, $\\vec{n}$ s.t. $n_k$ = the number of elements of class *k* in the sample, $\\sum_{k = 1}^K n_k = n$ (the number of Tweets from user *k* in the sample)\n",
    "* *__K__*: The total number of classes in the population, or the total number of users on Twitter talking about cats\n",
    " * now this of $\\vec{K}$ as the vetor with elements $K_j$ s.t. $K_j$ = the number of classes in the population with *j* elements (the number of users with *j* Tweets about cats)\n",
    "* $\\mathbf{\\vec{x}}$: $\\vec{x}$ s.t. $x_i$ = the total number of classes with *i* elements in the sample. $\\sum_{i = 1}^n x_i = $ the total number of classes in the sample; x is the observable statistic that we will use to estimate the number of classes in the population. This is a vector related to how many users have exactly *i* Tweets in the sample. One way to think about $\\vec{x}$ is that we are trying to estimate $x_0$, or the toal number of classes (users) who do not appear in the sample at all. \n",
    "* *__q__*: While I don't talk about this here, it is worth noting that Goodman's estimator can only be shown to be unbiased if *q < n*, where *q* is the total number of Tweets in the population (not the sample) from the most common class (most active user, if you like). I.e., *q* is the total number of Tweets generated by @CrazyCatLadyXxxX. In our case, we won't always know *q* exactly, but we can make an educated guess about whether or not it is smaller than the size of our sample.\n",
    "* *__S__*: The estimator for K. This is what we are solving for. $S_j$ is analogous to $K_j$, an estimator of the total number of classes with *j* elements in the population."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Section 2\n",
    "\n",
    "##### 2.1: Make the test populations (choose either Uniform or Poisson before proceeding)\n",
    "Feel free to change \"pop_size\" and \"percent_sample\" or mess with the number of classes/ distribution parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Imports\n",
    "from __future__ import division\n",
    "import numpy as np\n",
    "from scipy.misc import comb\n",
    "from numpy import random as np_random\n",
    "from collections import Counter\n",
    "from random import sample, choice\n",
    "from math import factorial, log, exp\n",
    "import fileinput\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# make a population with classes that are Poisson distributed\n",
    "l = 10\n",
    "pop_size = 10000\n",
    "population = list(np_random.poisson(lam = l, size = pop_size))\n",
    "percent_sample = 1\n",
    "\n",
    "print \"For a {}% sample of a Poisson({}) population of size {}: \\n\".format(percent_sample, l, pop_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# OR: make a population with classes that are Uniform distributed\n",
    "pop_size = 10000\n",
    "population = [choice(range(0,100)) for x in range(0,pop_size)]\n",
    "percent_sample = 1\n",
    "\n",
    "print \"For a {}% sample of a Uniform (integers 1 through 100) population of size {}: \\n\".format(percent_sample, l, pop_size)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Find $\\vec{x}$, look at the properties of the population and the sample"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# generate samples and print out info about the population\n",
    "print \"\\n The exact population demographics:\"\n",
    "pop_demographic = sorted(Counter(population).items(), key=lambda x: x[0])\n",
    "for key, value in pop_demographic:\n",
    "    print \"    {} items of class {}\".format(str(value).ljust(5), key)\n",
    "\n",
    "print \"\\n The demographics of the hypergeometric sample population:\"\n",
    "exactly_10_percent_sample = sample(population,int(pop_size*percent_sample/100))\n",
    "n = len(exactly_10_percent_sample)\n",
    "exactly_10_demographic_sample_only = Counter(exactly_10_percent_sample)\n",
    "exactly_10_demographic = Counter(exactly_10_percent_sample)\n",
    "for key in [x[0] for x in pop_demographic]:\n",
    "    if key not in exactly_10_demographic.keys():\n",
    "        exactly_10_demographic[key] = 0\n",
    "exactly_10_demographic = sorted(exactly_10_demographic.items(), key=lambda x: x[0])\n",
    "for key, value in exactly_10_demographic:\n",
    "    print \"    {} items of class {}\".format(str(value).ljust(5), key)\n",
    "print \" If the statistic x_{i} is the number of classes with i elements in the sample:\"\n",
    "x_counts_exact = sorted(Counter([x[1] for x in exactly_10_demographic]).items(), key=lambda x: x[0])\n",
    "for count, val in x_counts_exact:\n",
    "    print \"    {} classes appear {} times in the sample, x_{} = {}\".format(val, count, count, val)\n",
    "print \" So the observable vector x for this sample is \" \n",
    "print \" (recall that the number of classes appearing 0 times in the sample is not observable):\"\n",
    "x_vec = [0 for x in range(0,max([x[0] for x in x_counts_exact]))]\n",
    "for x in x_counts_exact:\n",
    "    if x[0] != 0:\n",
    "        x_vec[x[0]-1] = x[1]\n",
    "x_padded = np.pad(np.array(x_vec),(0,n-len(x_vec)),mode='constant',constant_values=0)\n",
    "print x_vec"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$\\vec{x}$ as a bar chart, a pdf. I've shown $x_0$ here in gray (because we know the exact composition of the population) even though $x_0$ is really unobservable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.bar(range(1,len(x_vec) + 1), x_vec, .45)\n",
    "ax.bar([0.05], [x_counts_exact[0][1]], .45, color = 'gray')\n",
    "ax.set_xticks([t + .22 for t in range(0,len(x_vec) + 1)])\n",
    "ax.set_xticklabels(range(0,len(x_vec) + 1))\n",
    "plt.ylabel(\"# of classes appearing x times in the sample\")\n",
    "plt.xlabel(\"x\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 2.2: Evaluating estimators\n",
    "\n",
    "In his paper, Goodman derives the estimator *S* (I won't go through the entire derivation here), except to note that *S* is the solution to the system of equations where:\n",
    "\n",
    "$$ x_i = \\sum_{j = 1}^{n} Pr(i \\; | \\; j,N,n)\\,S_j$$\n",
    "\n",
    "Note that the probabiltiy of having *i* elements of some class in the sample where that class has *j* elements in the total population is $$ Pr(i \\; | \\; j,N,n) = \\frac{ {j \\choose i} { N - j \\choose n - i} }{ {N \\choose n}} $$ which can be explained as: \n",
    "* $ { j \\choose i} $: from a class with *j* elements, choose *i* of them\n",
    "* $ { N-j \\choose n-i} $: from all other classes, choose *n-i* elements to make up an *n*-element sample\n",
    "* $ { N \\choose n} $: thte total number of ways to choose an *n*-element sample\n",
    "\n",
    "Then we multiply by $S_j$, the number of classes with *j* elements to get the expected value for $x_i$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# This would calculate the unbiased estimator by explicitly solving for S in the system of equations.\n",
    "# We don't do this because integer overflows. Not good. You can try it if you like for small populations/ samples\n",
    "'''\n",
    "eq_matrix = np.zeros((n,n))\n",
    "for i in xrange(0,n+1):\n",
    "    for j in xrange(0,n+1):\n",
    "        eq_matrix[i,j] = comb(i,j)*comb((pop_size - j), (n - i))/comb(pop_size, n)\n",
    "\n",
    "s_vec = np.linalg.solve(eq_matrix,x_padded)\n",
    "sum(s_vec)\n",
    "'''"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There's a closed form solution for $S_j$, presented in Goodman's work. Once again, I won't go through tthe derivation. I have modified the calculation to use logs to avoid integer overflows, but the calculation is the same."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# calculate the unbiased estimator S\n",
    "S = 0\n",
    "for i in xrange(1, n + 1):\n",
    "    x_i = x_padded[i - 1]\n",
    "    if x_i != 0:\n",
    "        sign = (-1)**(i+1)\n",
    "        log_numerator = sum([log(h) for h in range(pop_size - n, pop_size - n + i - 1 + 1)])\n",
    "        log_denominator = sum([log(k) for k in range(n - i, n - i + 1)])\n",
    "        # print x_i + sign*exp(log_numerator/log_denominator)*x_i\n",
    "        S += x_i + sign*exp(log_numerator/log_denominator)*x_i\n",
    "print \"The estimated number of total classes in the population is {}\".format(S)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 2.3\n",
    "\n",
    "As we have just observed, the unbiased estimator can have a very, very large error, to the point where it isn't even necessarily positive.\n",
    "\n",
    "Becuase of this, Goodman suggests some other biased estimators that have smaller error for some distributions of the population."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$ S' = N - \\frac{{N \\choose 2}}{{n \\choose 2}} x_2 $$\n",
    "\n",
    "gives a resonable lower bound when $(N)(x_2) \\leq n^2$ (not necessarily the case, see below)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# calculate the biased estimator S'\n",
    "S_1 = pop_size - ((pop_size)*(pop_size - 1))/((n)*(n - 1))*x_padded[2 - 1]\n",
    "print \"The estimator S' gives {} as the total number of classes in the population\".format(S_1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Or assume that the fraction $\\frac{n}{N}$ is equal to the fraction of classes that were sampled\n",
    "$$ S'' = \\frac{N}{n} \\sum{x} $$\n",
    "\n",
    "(not useful if we sampled most of the classes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# calculate S''\n",
    "S_2 = (pop_size/n) * sum(x_padded)\n",
    "print \"The estimator S'' gives {} as the total number of classes in the population\".format(S_2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Or assume that we sampled most of the classes (this is a hard lower bound for the total number of classes, there cannot be fewer classes than we have already observed):\n",
    "\n",
    "$$ \\sum_{i = 1}^n x_i $$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# calculate S'' (2)\n",
    "S_3 = sum(x_padded)\n",
    "print \"The estimator S'' gives {} as the total number of classes in the population\".format(S_3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Or, as is suggested in I.J. Good's paper, we can estimate $x_0$ as $x_1$ (assuming that there are the same number of unsampled classes as there are classes sampled exactly once, this holds for populations where there is not a large tail of very rare classes), and then estimate the proportion of classes that do appear in the sample as $\\left(1 - \\frac{x_1}{n}\\right)$\n",
    "$$ \\left(1 - \\frac{x_1}{n}\\right)^{-1} \\sum_{i = 1}^n x_i $$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# From I.J. Good's paper\n",
    "S_4 = (1/(1 - x_padded[0]/n))*(sum(x_padded))\n",
    "print \"The estimator suggested by I.J. Good gives {} as the number of classes in the population\".format(S_4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "print \"Reminder: the actual number of classes in the population is {}\".format(len(set(population)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Section 3\n",
    "\n",
    "##### 3.1: Load real Twitter data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "tweet_pop = []\n",
    "for line in fileinput.FileInput(\"craft_users_tweeting_soda_usernames.txt\"):\n",
    "    tweet_pop.append(line.strip())\n",
    "tweet_pop_size = len(tweet_pop)\n",
    "num_uniq_users_in_pop = len(set(tweet_pop))\n",
    "tweet_pop_demographic = sorted(Counter(tweet_pop).items(), key=lambda x: x[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "#print tweet_pop_demographic"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 3.2: Take a sample\n",
    "\n",
    "Feel free to change the sample size \"tweet_n\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "## Sample size, generate a new sample\n",
    "tweet_n = 3000\n",
    "sample_tweet_pop = sample(tweet_pop, tweet_n)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "sample_tweet_pop_demographic = Counter(sample_tweet_pop)\n",
    "for key in [x[0] for x in tweet_pop_demographic]:\n",
    "    if key not in sample_tweet_pop_demographic.keys():\n",
    "        sample_tweet_pop_demographic[key] = 0\n",
    "sample_tweet_pop_demographic = sorted(sample_tweet_pop_demographic.items(), key=lambda x: x[0])\n",
    "\n",
    "print \" If the statistic x_{i} is the number of classes with i elements in the sample:\"\n",
    "x_counts_tweets = sorted(Counter([x[1] for x in sample_tweet_pop_demographic]).items(), key=lambda x: x[0])\n",
    "for count, val in x_counts_tweets:\n",
    "    print \"    {} classes appear {} times in the sample, x_{} = {}\".format(val, count, count, val)\n",
    "print \" So the observable vector x for this sample is \" \n",
    "print \" (recall that the number of classes appearing 0 times in the sample is not observable):\"\n",
    "x_tweets_sample = [0 for x in range(0,max([x[0] for x in x_counts_tweets]))]\n",
    "for x in x_counts_tweets:\n",
    "    if x[0] != 0:\n",
    "        x_tweets_sample[x[0]-1] = x[1]\n",
    "x_tweets = np.pad(np.array(x_tweets_sample),(0,tweet_n-len(x_tweets_sample)),mode='constant',constant_values=0)\n",
    "print x_tweets_sample"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(10,6))\n",
    "ax.bar(range(0,len(x_counts_tweets)), [t[1] for t in x_counts_tweets], .45)\n",
    "#ax.bar([0.05], [x_counts_exact[0][1]], .45, color = 'gray')\n",
    "ax.set_xticks([t + .22 for t in range(0,len(x_tweets_sample) + 1)])\n",
    "ax.set_xticklabels(range(0,len(x_tweets_sample) + 1))\n",
    "plt.ylabel(\"# of classes appearing x times in the sample\")\n",
    "plt.xlabel(\"x (unobservable, because x_0, the number of unobserved classes, is included)\")\n",
    "plt.title = (\"Histogram of x\")\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(10,6))\n",
    "ax.bar(range(0,len(x_tweets_sample)), x_tweets_sample, .45)\n",
    "#ax.bar([0.05], [x_counts_exact[0][1]], .45, color = 'gray')\n",
    "ax.set_xticks([t + .22 for t in range(0,len(x_tweets_sample) + 1)])\n",
    "ax.set_xticklabels(range(1,len(x_tweets_sample) + 1))\n",
    "plt.ylabel(\"# of classes appearing x times in the sample\")\n",
    "plt.xlabel(\"x oberseved\")\n",
    "plt.title = (\"Histogram of x\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 3.2: Evaluating estimators (same content as 2.2, different data)\n",
    "\n",
    "In his paper, Goodman derives the estimator *S* (I won't go through the entire derivation here), except to note that *S* is the solution to the system of equations where:\n",
    "\n",
    "$$ x_i = \\sum_{j = 1}^{n} Pr(i \\; | \\; j,N,n)\\,S_j$$\n",
    "\n",
    "Note that the probabiltiy of having *i* elements of some class in the sample where that class has *j* elements in the total population is $$ Pr(i \\; | \\; j,N,n) = \\frac{ {j \\choose i} { N - j \\choose n - i} }{ {N \\choose n}} $$ which can be explained as: \n",
    "* $ { j \\choose i} $: from a class with *j* elements, choose *i* of them\n",
    "* $ { N-j \\choose n-i} $: from all other classes, choose *n-i* elements to make up an *n*-element sample\n",
    "* $ { N \\choose n} $: thte total number of ways to choose an *n*-element sample\n",
    "\n",
    "Then we multiply by $S_j$, the number of classes with *j* elements to get the expected value for $x_i$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# calculate the unbiased estimator S\n",
    "tweet_S = 0\n",
    "test = []\n",
    "for i in xrange(1, tweet_n + 1):\n",
    "    x_i = x_tweets[i - 1]\n",
    "    if x_i != 0:\n",
    "        sign = (-1)**(i+1)\n",
    "        log_numerator = sum([log(h) for h in range(tweet_pop_size - tweet_n, tweet_pop_size - tweet_n + i - 1 + 1)])\n",
    "        log_denominator = sum([log(k) for k in range(tweet_n - i, tweet_n - i + 1)])\n",
    "        test.append(x_i + sign*exp(log_numerator/log_denominator)*x_i)\n",
    "        tweet_S += x_i + sign*exp(log_numerator/log_denominator)*x_i\n",
    "print \"The estimated number of total classes in the population is {}\".format(tweet_S)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$ S' = N - \\frac{{N \\choose 2}}{{n \\choose 2}} x_2 $$\n",
    "\n",
    "gives a resonable lower bound when $(N)(x_2) \\leq n^2$ (not necessarily the case, see below)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# calculate the biased estimator S'\n",
    "tweet_S_1 = tweet_pop_size - ((tweet_pop_size)*(tweet_pop_size - 1))/((tweet_n)*(tweet_n - 1))*x_tweets[2 - 1]\n",
    "print \"The estimator S' gives {} as the total number of classes in the population\".format(tweet_S_1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Or assume that the fraction $\\frac{n}{N}$ is equal to the fraction of classes that were sampled\n",
    "$$ S'' = \\frac{N}{n} \\sum{x} $$\n",
    "\n",
    "(not useful if we sampled most of the classes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# calculate S''\n",
    "tweet_S_2 = (tweet_pop_size/tweet_n) * sum(x_tweets)\n",
    "print \"The estimator S'' gives {} as the total number of classes in the population\".format(tweet_S_2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Or assume that we sampled most of the classes (this is a hard lower bound for the total number of classes, there cannot be fewer classes than we have already observed):\n",
    "\n",
    "$$ \\sum_{i = 1}^n x_i $$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# calculate S'' (2)\n",
    "tweet_S_3 = sum(x_tweets)\n",
    "print \"The estimator S'' gives {} as the total number of classes in the population\".format(tweet_S_3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Or, as is suggested in I.J. Good's paper, we can estimate $x_0$ as $x_1$ (assuming that there are the same number of unsampled classes as there are classes sampled exactly once, this holds for populations where there is not a large tail of very rare classes), and then estimate the proportion of classes that do appear in the sample as $\\left(1 - \\frac{x_1}{n}\\right)$\n",
    "$$ \\left(1 - \\frac{x_1}{n}\\right)^{-1} \\sum_{i = 1}^n x_i $$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# From I.J. Good's paper\n",
    "tweet_S_4 = (1/(1 - x_tweets[0]/tweet_n))*(sum(x_tweets))\n",
    "print \"The estimator suggested by I.J. Good gives {} as the number of classes in the population\".format(tweet_S_4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "print \"Reminder: the actual number of classes (unique users) in the population is: {}\".format(num_uniq_users_in_pop)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 3.3: Future work\n",
    "\n",
    "Disclaimer: it's completely possible that none of this makes any sense. Just poking around.\n",
    "\n",
    "Try cutting off $\\vec{x}$ at the first zero (to remove a long tail) and treating very common classes (outliers) differently from the less common ones."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "first_cutoff = np.where(x_tweets<=0)[0][0]\n",
    "x_tweets_ = x_tweets[0:first_cutoff]\n",
    "outliers = []\n",
    "for i,k in enumerate(x_tweets[first_cutoff:]):\n",
    "    if k != 0:\n",
    "        outliers.append((i + first_cutoff + 1, k))\n",
    "# assume that these more common classes were sampled pretty evenly\n",
    "num_outliers_in_sample = sum([x[0]*x[1] for x in outliers])\n",
    "num_outliers_in_pop = sum([x[0]*x[1]*tweet_pop_size/tweet_n for x in outliers])\n",
    "num_outlier_classes = sum([x[1] for x in outliers])\n",
    "\n",
    "tweet_pop_size_ = int(tweet_pop_size - num_outliers_in_pop)\n",
    "tweet_n_ = int(tweet_n - num_outliers_in_sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# calculate the unbiased estimator S\n",
    "tweet_S = 0\n",
    "for i in xrange(1, len(x_tweets_) + 1):\n",
    "    x_i = x_tweets_[i - 1]\n",
    "    sign = (-1)**(i+1)\n",
    "    log_numerator = sum([log(h) for h in range(tweet_pop_size_ - tweet_n_, tweet_pop_size_ - tweet_n_ + i - 1 + 1)])\n",
    "    log_denominator = sum([log(k) for k in range(tweet_n_ - i, tweet_n_ - i + 1)])\n",
    "    # print x_i + sign*exp(log_numerator/log_denominator)*x_i\n",
    "    tweet_S += x_i + sign*exp(log_numerator/log_denominator)*x_i\n",
    "print \"The estimated number of total classes in the population is {}\".format(int(tweet_S + num_outlier_classes))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Look at the difference in frequency distribution between classes that do appear in the sample and classes that don't."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Analysis of the classes that *do not* appear in the sample\n",
    "users_not_in_sample = set([x[0] for x in sample_tweet_pop_demographic if x[1]==0])\n",
    "users_not_in_sample_demographic = []\n",
    "users_in_sample_demographic = []\n",
    "for u in tweet_pop_demographic:\n",
    "    if u[0] in users_not_in_sample:\n",
    "        users_not_in_sample_demographic.append(u)\n",
    "    else:\n",
    "        users_in_sample_demographic.append(u)\n",
    "counts_users_in_sample = Counter([x[1] for x in users_in_sample_demographic]).items()\n",
    "counts_users_not_in_sample = Counter([x[1] for x in users_not_in_sample_demographic]).items()\n",
    "\n",
    "x_users_in_sample = [0 for x in range(0,max([x[0] for x in counts_users_in_sample]))]\n",
    "for x in counts_users_in_sample:\n",
    "    if x[0] != 0:\n",
    "        x_users_in_sample[x[0]-1] = x[1]\n",
    "\n",
    "x_users_not_in_sample = [0 for x in range(0,max([x[0] for x in counts_users_not_in_sample]))]\n",
    "for x in counts_users_not_in_sample:\n",
    "    if x[0] != 0:\n",
    "        x_users_not_in_sample[x[0]-1] = x[1]\n",
    "\n",
    "print x_users_in_sample\n",
    "print x_users_not_in_sample"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(10,6))\n",
    "ax.bar(range(0,len(x_users_not_in_sample)), x_users_not_in_sample, .45)\n",
    "ax.set_xticks([t + .22 for t in range(0,len(x_users_not_in_sample) + 1)])\n",
    "ax.set_xticklabels(range(0,len(x_users_not_in_sample) + 1))\n",
    "plt.ylabel(\"# of classes appearing x times in the pop and at least once in the sample\")\n",
    "plt.xlabel(\"x \")\n",
    "plt.title = (\"Histogram of x\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "x_users_not_in_sample[0]/sum(x_users_not_in_sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "x_users_in_sample[0]/sum(x_users_in_sample)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "##### This seems to give a resonable estimate occasionally\n",
    "\n",
    "Estimate $x_0$ as an extrapolation of the ratio between $x_1$ and $x_2$ (as opposed to simply estimating $x_0 = x_1$)\n",
    "\n",
    "$$ \\frac{x_1}{x_2} = \\frac{x_0}{x_1} $$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "print \"The total number of classes in the population: {}\".format(int((x_tweets[0]*x_tweets[0]/x_tweets[1]) + sum(x_tweets)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "print \"Reminder: the actual number of classes (unique users) in the population is: {}\".format(num_uniq_users_in_pop)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "test_estimator = []\n",
    "test_estimator_2 = []\n",
    "test_lower_bound = []\n",
    "test_upper_bound = []\n",
    "test_goodman_s1 = []\n",
    "test_goodman_s2 = []\n",
    "test_ijgood = []\n",
    "for test_tweet_n in range(400,int(38000/2),100):\n",
    "    for t in range(0,10):\n",
    "        ## Sample size, generate a new sample\n",
    "        test_sample_tweet_pop = sample(tweet_pop, test_tweet_n)\n",
    "\n",
    "        test_sample_tweet_pop_demographic = sorted(Counter(test_sample_tweet_pop).items(), key=lambda x: x[0])\n",
    "\n",
    "        test_x_counts_tweets = sorted(Counter([x[1] for x in test_sample_tweet_pop_demographic]).items(), key=lambda x: x[0])\n",
    "\n",
    "        test_x_tweets_sample = [0 for x in range(0,max([x[0] for x in test_x_counts_tweets]))]\n",
    "        for x in test_x_counts_tweets:\n",
    "            if x[0] != 0:\n",
    "                test_x_tweets_sample[x[0]-1] = x[1]\n",
    "        test_x_tweets = np.pad(np.array(test_x_tweets_sample),(0,test_tweet_n-len(test_x_tweets_sample)),mode='constant',constant_values=0)\n",
    "\n",
    "        sum_test_x_tweets = sum(test_x_tweets)\n",
    "        \n",
    "        est_1 = ((test_x_tweets[0]*test_x_tweets[0]/test_x_tweets[1]) + sum_test_x_tweets)\n",
    "        est_2 = (1/(1 - test_x_tweets[0]/test_tweet_n))*(sum_test_x_tweets)\n",
    "        est_3 = sum_test_x_tweets*(len(tweet_pop)/test_tweet_n)\n",
    "        est_4 = len(tweet_pop) - ((len(tweet_pop))*(len(tweet_pop) - 1))/((test_tweet_n)*(test_tweet_n - 1))*test_x_tweets[2 - 1]\n",
    "        \n",
    "        test_estimator.append((test_tweet_n, est_1/num_uniq_users_in_pop))\n",
    "        test_ijgood.append((test_tweet_n, est_2/num_uniq_users_in_pop))\n",
    "        test_goodman_s2.append((test_tweet_n, est_3/num_uniq_users_in_pop))\n",
    "        test_goodman_s1.append((test_tweet_n, est_4/num_uniq_users_in_pop))\n",
    "        test_lower_bound.append((test_tweet_n, sum_test_x_tweets/num_uniq_users_in_pop))\n",
    "        test_upper_bound.append((test_tweet_n, (len(tweet_pop) - test_tweet_n + sum_test_x_tweets)/num_uniq_users_in_pop))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(10,6))\n",
    "ax.set_ylim([-.5,2.25])\n",
    "plt.plot([x[0]/len(tweet_pop)*100 for x in test_estimator], [x[1] for x in test_estimator],'r.')\n",
    "plt.plot([x[0]/len(tweet_pop)*100 for x in test_ijgood], [x[1] for x in test_ijgood],'b.')\n",
    "plt.plot([x[0]/len(tweet_pop)*100 for x in test_goodman_s2], [x[1] for x in test_goodman_s2],'g.')\n",
    "plt.plot([x[0]/len(tweet_pop)*100 for x in test_goodman_s1], [x[1] for x in test_goodman_s1],'m.')\n",
    "plt.plot([x[0]/len(tweet_pop)*100 for x in test_lower_bound], [x[1] for x in test_lower_bound],'k')\n",
    "plt.plot([x[0]/len(tweet_pop)*100 for x in test_upper_bound], [x[1] for x in test_upper_bound],'k')\n",
    "plt.plot([x[0]/len(tweet_pop)*100 for x in test_upper_bound], [1 for x in test_upper_bound],'k-')\n",
    "plt.xlabel(\"Sample is x% of the total pop\")\n",
    "plt.ylabel(\"Fraction of the true number of classes\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
